"""
策略回测数据库记录器
用于记录策略回测结果和交易日志
"""

import sqlite3
import json
import pandas as pd
from datetime import datetime
from typing import Dict, Any, Optional, List


class StrategyDBRecorder:
    """策略回测数据库记录器"""
    
    def __init__(self, db_path: str = "trading_backtest.db"):
        self.db_path = db_path
        self.init_database()
    
    def init_database(self):
        """初始化数据库表结构"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 创建策略回测结果表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS backtest_results (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                strategy_name TEXT NOT NULL,
                symbol TEXT NOT NULL,
                start_date TEXT NOT NULL,
                end_date TEXT NOT NULL,
                init_balance REAL NOT NULL,
                final_balance REAL,
                total_return REAL,
                annual_yield REAL,
                sharpe_ratio REAL,
                max_drawdown REAL,
                winning_rate REAL,
                profit_loss_ratio REAL,
                open_times INTEGER,
                close_times INTEGER,
                run_duration_seconds REAL,
                strategy_params TEXT,  -- JSON格式存储参数
                test_timestamp TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建每日交易汇总表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS daily_trade_summary (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                backtest_id INTEGER,
                trade_date TEXT NOT NULL,
                trades_count INTEGER DEFAULT 0,
                positions_count INTEGER DEFAULT 0,
                daily_pnl REAL DEFAULT 0,
                trades_data TEXT,  -- JSON格式存储当日所有交易
                positions_data TEXT,  -- JSON格式存储当日持仓信息
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (backtest_id) REFERENCES backtest_results (id)
            )
        ''')
        
        # 创建详细交易记录表（从trades中解析出的单笔交易）
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS trade_details (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                backtest_id INTEGER,
                daily_summary_id INTEGER,
                trade_date TEXT NOT NULL,
                datetime TEXT,
                symbol TEXT,
                direction TEXT,  -- LONG/SHORT
                offset TEXT,     -- OPEN/CLOSE
                volume INTEGER,
                price REAL,
                commission REAL,
                trade_id TEXT,
                order_id TEXT,
                profit_loss REAL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (backtest_id) REFERENCES backtest_results (id),
                FOREIGN KEY (daily_summary_id) REFERENCES daily_trade_summary (id)
            )
        ''')
        
        # 创建持仓记录表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS position_records (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                backtest_id INTEGER,
                daily_summary_id INTEGER,
                trade_date TEXT NOT NULL,
                symbol TEXT,
                position_long INTEGER DEFAULT 0,
                position_short INTEGER DEFAULT 0,
                volume_long INTEGER DEFAULT 0,
                volume_short INTEGER DEFAULT 0,
                open_price_long REAL DEFAULT 0,
                open_price_short REAL DEFAULT 0,
                float_profit REAL DEFAULT 0,
                position_profit REAL DEFAULT 0,
                margin REAL DEFAULT 0,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (backtest_id) REFERENCES backtest_results (id),
                FOREIGN KEY (daily_summary_id) REFERENCES daily_trade_summary (id)
            )
        ''')
        
        # 创建账户记录表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS account_records (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                backtest_id INTEGER,
                daily_summary_id INTEGER,
                trade_date TEXT NOT NULL,
                currency TEXT DEFAULT 'CNY',
                pre_balance REAL DEFAULT 0,
                static_balance REAL DEFAULT 0,
                balance REAL DEFAULT 0,
                available REAL DEFAULT 0,
                float_profit REAL DEFAULT 0,
                position_profit REAL DEFAULT 0,
                close_profit REAL DEFAULT 0,
                frozen_margin REAL DEFAULT 0,
                margin REAL DEFAULT 0,
                frozen_commission REAL DEFAULT 0,
                commission REAL DEFAULT 0,
                frozen_premium REAL DEFAULT 0,
                premium REAL DEFAULT 0,
                deposit REAL DEFAULT 0,
                withdraw REAL DEFAULT 0,
                risk_ratio REAL DEFAULT 0,
                market_value REAL DEFAULT 0,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (backtest_id) REFERENCES backtest_results (id),
                FOREIGN KEY (daily_summary_id) REFERENCES daily_trade_summary (id)
            )
        ''')
        
        conn.commit()
        conn.close()
    
    def save_backtest_result(self, 
                           strategy_name: str,
                           symbol: str,
                           start_date: str,
                           end_date: str,
                           init_balance: float,
                           final_balance: float,
                           tqsdk_stats: Any,
                           run_duration: float,
                           strategy_params: Dict,
                           test_timestamp: str = None) -> int:
        """
        保存回测结果到数据库
        
        Returns:
            int: 插入记录的ID
        """
        if test_timestamp is None:
            test_timestamp = datetime.now().isoformat()
        
        # 从tqsdk_stats提取统计信息
        total_return = 0
        annual_yield = 0
        sharpe_ratio = 0
        max_drawdown = 0
        winning_rate = 0
        profit_loss_ratio = 0
        open_times = 0
        close_times = 0
        
        if tqsdk_stats:
            total_return = getattr(tqsdk_stats, 'ror', 0)
            annual_yield = getattr(tqsdk_stats, 'annual_yield', 0)
            sharpe_ratio = getattr(tqsdk_stats, 'sharpe_ratio', 0)
            max_drawdown = getattr(tqsdk_stats, 'max_drawdown', 0)
            winning_rate = getattr(tqsdk_stats, 'winning_rate', 0)
            profit_loss_ratio = getattr(tqsdk_stats, 'profit_loss_ratio', 0)
            open_times = getattr(tqsdk_stats, 'open_times', 0)
            close_times = getattr(tqsdk_stats, 'close_times', 0)
        
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT INTO backtest_results (
                strategy_name, symbol, start_date, end_date, init_balance, final_balance,
                total_return, annual_yield, sharpe_ratio, max_drawdown, winning_rate,
                profit_loss_ratio, open_times, close_times, run_duration_seconds,
                strategy_params, test_timestamp
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            strategy_name, symbol, start_date, end_date, init_balance, final_balance,
            total_return, annual_yield, sharpe_ratio, max_drawdown, winning_rate,
            profit_loss_ratio, open_times, close_times, run_duration,
            json.dumps(strategy_params, ensure_ascii=False), test_timestamp
        ))
        
        backtest_id = cursor.lastrowid
        conn.commit()
        conn.close()
        
        return backtest_id
    
    def save_trade_logs(self, backtest_id: int, trade_log: Any):
        """
        保存交易日志到数据库
        处理TqSDK的真实日期字典格式: {'2023-01-04': {'trades': [trade_obj], 'account': account_obj, 'positions': {symbol: pos_obj}}}
        
        Args:
            backtest_id: 回测结果ID
            trade_log: TqSDK交易日志对象（日期字典格式）
        
        Returns:
            int: 保存的记录数量
        """
        if not trade_log:
            return 0
        
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        total_records = 0
        
        try:
            # 处理日期字典格式的trade_log
            if isinstance(trade_log, dict):
                for trade_date, date_data in trade_log.items():
                    if not isinstance(date_data, dict):
                        continue
                    
                    # 提取当日数据
                    trades = date_data.get('trades', [])  # 交易记录列表
                    account = date_data.get('account', {})  # 账户信息
                    positions = date_data.get('positions', {})  # 持仓信息字典
                    
                    # 计算当日统计
                    trades_count = len(trades) if trades else 0
                    positions_count = len(positions) if positions else 0
                    
                    # 从账户信息中获取当日盈亏
                    daily_pnl = 0
                    if account:
                        if isinstance(account, dict):
                            daily_pnl = account.get('close_profit', 0)
                        else:
                            daily_pnl = getattr(account, 'close_profit', 0)
                    
                    # 保存每日汇总
                    cursor.execute('''
                        INSERT INTO daily_trade_summary (
                            backtest_id, trade_date, trades_count, positions_count, daily_pnl,
                            trades_data, positions_data
                        ) VALUES (?, ?, ?, ?, ?, ?, ?)
                    ''', (
                        backtest_id,
                        trade_date,
                        trades_count,
                        positions_count,
                        daily_pnl,
                        json.dumps(self._serialize_trades_list(trades), ensure_ascii=False),
                        json.dumps(self._serialize_positions_dict(positions), ensure_ascii=False)
                    ))
                    
                    daily_summary_id = cursor.lastrowid
                    
                    # 保存详细交易记录
                    if trades:
                        trade_details = self._extract_trade_details_from_list(trades, backtest_id, daily_summary_id, trade_date)
                        if trade_details:
                            cursor.executemany('''
                                INSERT INTO trade_details (
                                    backtest_id, daily_summary_id, trade_date, datetime, symbol,
                                    direction, offset, volume, price, commission, trade_id, order_id, profit_loss
                                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                            ''', trade_details)
                    
                    # 保存持仓记录
                    if positions:
                        position_records = self._extract_position_records_from_dict(positions, backtest_id, daily_summary_id, trade_date)
                        if position_records:
                            cursor.executemany('''
                                INSERT INTO position_records (
                                    backtest_id, daily_summary_id, trade_date, symbol,
                                    position_long, position_short, volume_long, volume_short,
                                    open_price_long, open_price_short, float_profit, position_profit, margin
                                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                            ''', position_records)
                    
                    # 保存账户记录
                    if account:
                        account_record = self._extract_account_record(account, backtest_id, daily_summary_id, trade_date)
                        if account_record:
                            cursor.execute('''
                                INSERT INTO account_records (
                                    backtest_id, daily_summary_id, trade_date, currency,
                                    pre_balance, static_balance, balance, available, float_profit, position_profit,
                                    close_profit, frozen_margin, margin, frozen_commission, commission,
                                    frozen_premium, premium, deposit, withdraw, risk_ratio, market_value
                                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                            ''', account_record)
                    
                    total_records += 1
            
            conn.commit()
            
        except Exception as e:
            print(f"❌ 保存交易日志失败: {e}")
            import traceback
            traceback.print_exc()
            conn.rollback()
        finally:
            conn.close()
        
        return total_records
    
    def _serialize_trades_list(self, trades):
        """序列化交易列表为可JSON化的格式"""
        if not trades:
            return []
        
        serialized = []
        
        for trade in trades:
            serialized.append(self._serialize_single_trade(trade))
        
        return serialized
    
    def _serialize_single_trade(self, trade):
        """序列化单个交易对象"""
        if hasattr(trade, '__dict__'):
            # 对象属性
            return {
                attr: getattr(trade, attr, None)
                for attr in dir(trade)
                if not attr.startswith('_') and not callable(getattr(trade, attr, None))
            }
        elif isinstance(trade, dict):
            return trade
        else:
            return str(trade)
    
    def _serialize_positions_dict(self, positions):
        """序列化持仓字典为可JSON化的格式"""
        if not positions:
            return {}
        
        serialized = {}
        
        for symbol, position in positions.items():
            serialized[symbol] = self._serialize_single_position(position)
        
        return serialized
    
    def _serialize_single_position(self, position):
        """序列化单个持仓对象"""
        if hasattr(position, '__dict__'):
            # 对象属性
            return {
                attr: getattr(position, attr, None)
                for attr in dir(position)
                if not attr.startswith('_') and not callable(getattr(position, attr, None))
            }
        elif isinstance(position, dict):
            return position
        else:
            return str(position)
    
    def _extract_trade_details_from_list(self, trades, backtest_id, daily_summary_id, trade_date):
        """从交易列表中提取详细交易记录"""
        trade_details = []
        
        for trade in trades:
            # 转换时间戳为可读格式
            datetime_str = self._convert_timestamp(self._get_trade_attr(trade, 'trade_date_time', 0))
            
            detail = [
                backtest_id,
                daily_summary_id,
                trade_date,
                datetime_str,
                self._get_trade_attr(trade, 'instrument_id', ''),  # TqSDK使用instrument_id
                self._get_trade_attr(trade, 'direction', ''),
                self._get_trade_attr(trade, 'offset', ''),
                int(self._get_trade_attr(trade, 'volume', 0)),
                float(self._get_trade_attr(trade, 'price', 0)),
                float(self._get_trade_attr(trade, 'commission', 0)),
                self._get_trade_attr(trade, 'trade_id', ''),
                self._get_trade_attr(trade, 'order_id', ''),
                0  # TqSDK交易记录中没有直接的profit_loss字段
            ]
            trade_details.append(detail)
        
        return trade_details
    
    def _extract_position_records_from_dict(self, positions, backtest_id, daily_summary_id, trade_date):
        """从持仓字典中提取持仓记录"""
        position_records = []
        
        for symbol, position in positions.items():
            # 处理NaN值
            open_price_long = self._get_trade_attr(position, 'open_price_long', 0)
            open_price_short = self._get_trade_attr(position, 'open_price_short', 0)
            
            # 将NaN转换为0
            if str(open_price_long).lower() == 'nan':
                open_price_long = 0
            if str(open_price_short).lower() == 'nan':
                open_price_short = 0
            
            record = [
                backtest_id,
                daily_summary_id,
                trade_date,
                self._get_trade_attr(position, 'instrument_id', symbol),  # 使用instrument_id或symbol
                int(self._get_trade_attr(position, 'pos_long', 0)),  # TqSDK使用pos_long
                int(self._get_trade_attr(position, 'pos_short', 0)),  # TqSDK使用pos_short
                int(self._get_trade_attr(position, 'volume_long', 0)),
                int(self._get_trade_attr(position, 'volume_short', 0)),
                float(open_price_long),
                float(open_price_short),
                float(self._get_trade_attr(position, 'float_profit', 0)),
                float(self._get_trade_attr(position, 'position_profit', 0)),
                float(self._get_trade_attr(position, 'margin', 0))
            ]
            position_records.append(record)
        
        return position_records
    
    def _extract_account_record(self, account, backtest_id, daily_summary_id, trade_date):
        """从账户信息中提取账户记录"""
        if not account:
            return None
        
        # 处理NaN值的函数
        def safe_float(value):
            if str(value).lower() == 'nan':
                return 0.0
            try:
                return float(value)
            except:
                return 0.0
        
        record = [
            backtest_id,
            daily_summary_id,
            trade_date,
            self._get_trade_attr(account, 'currency', 'CNY'),
            safe_float(self._get_trade_attr(account, 'pre_balance', 0)),
            safe_float(self._get_trade_attr(account, 'static_balance', 0)),
            safe_float(self._get_trade_attr(account, 'balance', 0)),
            safe_float(self._get_trade_attr(account, 'available', 0)),
            safe_float(self._get_trade_attr(account, 'float_profit', 0)),
            safe_float(self._get_trade_attr(account, 'position_profit', 0)),
            safe_float(self._get_trade_attr(account, 'close_profit', 0)),
            safe_float(self._get_trade_attr(account, 'frozen_margin', 0)),
            safe_float(self._get_trade_attr(account, 'margin', 0)),
            safe_float(self._get_trade_attr(account, 'frozen_commission', 0)),
            safe_float(self._get_trade_attr(account, 'commission', 0)),
            safe_float(self._get_trade_attr(account, 'frozen_premium', 0)),
            safe_float(self._get_trade_attr(account, 'premium', 0)),
            safe_float(self._get_trade_attr(account, 'deposit', 0)),
            safe_float(self._get_trade_attr(account, 'withdraw', 0)),
            safe_float(self._get_trade_attr(account, 'risk_ratio', 0)),
            safe_float(self._get_trade_attr(account, 'market_value', 0))
        ]
        
        return record
    
    def _get_trade_attr(self, obj, attr_name, default_value):
        """安全获取交易对象的属性值"""
        if hasattr(obj, attr_name):
            return getattr(obj, attr_name, default_value)
        elif isinstance(obj, dict):
            return obj.get(attr_name, default_value)
        else:
            return default_value
    
    def _convert_timestamp(self, timestamp):
        """将纳秒时间戳转换为可读格式"""
        if not timestamp or timestamp == 0:
            return ''
        
        try:
            # TqSDK使用纳秒时间戳，需要转换为秒
            timestamp_seconds = timestamp / 1_000_000_000
            dt = datetime.fromtimestamp(timestamp_seconds)
            return dt.strftime('%Y-%m-%d %H:%M:%S')
        except:
            return str(timestamp)
    
    def get_backtest_results(self, limit: int = 10) -> pd.DataFrame:
        """获取回测结果"""
        conn = sqlite3.connect(self.db_path)
        
        query = '''
            SELECT * FROM backtest_results 
            ORDER BY created_at DESC 
            LIMIT ?
        '''
        
        df = pd.read_sql_query(query, conn, params=(limit,))
        conn.close()
        
        return df
    
    def get_daily_trade_summary(self, backtest_id: int) -> pd.DataFrame:
        """获取指定回测的每日交易汇总"""
        conn = sqlite3.connect(self.db_path)
        
        query = '''
            SELECT * FROM daily_trade_summary 
            WHERE backtest_id = ?
            ORDER BY trade_date
        '''
        
        df = pd.read_sql_query(query, conn, params=(backtest_id,))
        conn.close()
        
        return df
    
    def get_trade_details(self, backtest_id: int) -> pd.DataFrame:
        """获取指定回测的详细交易记录"""
        conn = sqlite3.connect(self.db_path)
        
        query = '''
            SELECT * FROM trade_details 
            WHERE backtest_id = ?
            ORDER BY trade_date, datetime
        '''
        
        df = pd.read_sql_query(query, conn, params=(backtest_id,))
        conn.close()
        
        return df
    
    def get_position_records(self, backtest_id: int) -> pd.DataFrame:
        """获取指定回测的持仓记录"""
        conn = sqlite3.connect(self.db_path)
        
        query = '''
            SELECT * FROM position_records 
            WHERE backtest_id = ?
            ORDER BY trade_date
        '''
        
        df = pd.read_sql_query(query, conn, params=(backtest_id,))
        conn.close()
        
        return df
    
    def get_account_records(self, backtest_id: int) -> pd.DataFrame:
        """获取指定回测的账户记录"""
        conn = sqlite3.connect(self.db_path)
        
        query = '''
            SELECT * FROM account_records 
            WHERE backtest_id = ?
            ORDER BY trade_date
        '''
        
        df = pd.read_sql_query(query, conn, params=(backtest_id,))
        conn.close()
        
        return df
    
    def get_trade_logs(self, backtest_id: int) -> pd.DataFrame:
        """获取指定回测的交易日志（兼容旧接口）"""
        return self.get_trade_details(backtest_id)
    
    def get_strategy_summary(self) -> pd.DataFrame:
        """获取策略汇总统计"""
        conn = sqlite3.connect(self.db_path)
        
        query = '''
            SELECT 
                strategy_name,
                COUNT(*) as test_count,
                AVG(total_return) as avg_return,
                MAX(total_return) as best_return,
                MIN(total_return) as worst_return,
                AVG(sharpe_ratio) as avg_sharpe,
                AVG(max_drawdown) as avg_drawdown,
                AVG(winning_rate) as avg_win_rate,
                SUM(open_times) as total_trades
            FROM backtest_results 
            GROUP BY strategy_name
            ORDER BY avg_return DESC
        '''
        
        df = pd.read_sql_query(query, conn)
        conn.close()
        
        return df
    
    def export_results(self, backtest_id: int, export_path: str = None) -> str:
        """导出回测结果和交易日志"""
        if export_path is None:
            export_path = f"backtest_export_{backtest_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"
        
        # 获取回测结果
        conn = sqlite3.connect(self.db_path)
        
        result_query = 'SELECT * FROM backtest_results WHERE id = ?'
        result_df = pd.read_sql_query(result_query, conn, params=(backtest_id,))
        
        # 获取各种交易数据
        daily_summary_query = 'SELECT * FROM daily_trade_summary WHERE backtest_id = ?'
        daily_summary_df = pd.read_sql_query(daily_summary_query, conn, params=(backtest_id,))
        
        trade_details_query = 'SELECT * FROM trade_details WHERE backtest_id = ?'
        trade_details_df = pd.read_sql_query(trade_details_query, conn, params=(backtest_id,))
        
        position_records_query = 'SELECT * FROM position_records WHERE backtest_id = ?'
        position_records_df = pd.read_sql_query(position_records_query, conn, params=(backtest_id,))
        
        account_records_query = 'SELECT * FROM account_records WHERE backtest_id = ?'
        account_records_df = pd.read_sql_query(account_records_query, conn, params=(backtest_id,))
        
        conn.close()
        
        # 导出到Excel
        with pd.ExcelWriter(export_path, engine='openpyxl') as writer:
            result_df.to_excel(writer, sheet_name='回测结果', index=False)
            daily_summary_df.to_excel(writer, sheet_name='每日汇总', index=False)
            trade_details_df.to_excel(writer, sheet_name='交易明细', index=False)
            position_records_df.to_excel(writer, sheet_name='持仓记录', index=False)
            account_records_df.to_excel(writer, sheet_name='账户记录', index=False)
        
        return export_path
    
    def cleanup_old_records(self, keep_days: int = 30) -> int:
        """清理旧记录"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 获取要删除的记录ID
        cursor.execute('''
            SELECT id FROM backtest_results 
            WHERE created_at < datetime('now', '-{} days')
        '''.format(keep_days))
        
        old_records = cursor.fetchall()
        
        if old_records:
            old_ids = [r[0] for r in old_records]
            
            # 按照外键依赖顺序删除
            # 1. 删除交易明细
            cursor.executemany('DELETE FROM trade_details WHERE backtest_id = ?', 
                             [(id,) for id in old_ids])
            
            # 2. 删除持仓记录
            cursor.executemany('DELETE FROM position_records WHERE backtest_id = ?', 
                             [(id,) for id in old_ids])
            
            # 3. 删除账户记录
            cursor.executemany('DELETE FROM account_records WHERE backtest_id = ?', 
                             [(id,) for id in old_ids])
            
            # 4. 删除每日汇总
            cursor.executemany('DELETE FROM daily_trade_summary WHERE backtest_id = ?', 
                             [(id,) for id in old_ids])
            
            # 4. 删除回测结果
            cursor.executemany('DELETE FROM backtest_results WHERE id = ?', 
                             [(id,) for id in old_ids])
            
            conn.commit()
        
        conn.close()
        
        return len(old_records) if old_records else 0